import sys
import csv
from collections import defaultdict
import matplotlib.pyplot as plt
import re
import argparse

csv.field_size_limit(500 * 1024 * 1024)


def plt_memory(memory_changes, mem_type, need_cumlate=True):
    if not memory_changes:
        return
    sorted_times = sorted(memory_changes.keys())
    cumulative_memory = 0
    times = []
    memory_usage = []

    for time in sorted_times:
        if need_cumlate:
            cumulative_memory += memory_changes[time]
            times.append(time)
            memory_usage.append(cumulative_memory)
        else:
            times.append(time)
            memory_usage.append(memory_changes[time])

    max_index = memory_usage.index(max(memory_usage))
    peak_time = times[max_index]
    peak_memory_usage = memory_usage[max_index]

    if mem_type == 'memory_without_fragment':
        label = 'memory_without_fragment'
        color = 'b'
        height = 10
    elif mem_type == 'theory':
        label = 'theory_memory'
        color = 'r'
        height = 50
    else:
        label = 'memory_with_fragment'
        color = 'y'
        height = 50
    

    plt.plot(times[:-2], memory_usage[:-2], marker='o', linestyle='-', color=color, label=label)
    plt.annotate(f'(time: {peak_time}, memory(bytes): {peak_memory_usage})',  # Annotation text, the peak's coordinates
                 (peak_time, peak_memory_usage),  # The point to annotate
                 textcoords="offset points",  # How to position the text
                 xytext=(0, height),  # Distance from text to points (x,y)
                 ha='center',  # Horizontal alignment can be left, right or center
                 color=color)
    print(f'{mem_type} peak time: {peak_time}, peak memory usage: {peak_memory_usage} bytes')

    return peak_time


def tracker(filename, boost):
    objects = []
    real_memory_changes = defaultdict(int)
    theory_memory_changes = defaultdict(int)
    actual_memory = defaultdict(int)
    with open(filename, mode='r', encoding='utf-8') as file:
        # 使用 DictReader 读取文件，自动将每一行转换为一个字典
        reader = csv.DictReader(file)
        # real memory
        for row in reader:
            objects.append(row)
            start = int(row['start_time_stamp'])
            end = int(row['end_time_stamp'])
            size = int(row['size'])
            if 'actual_peak_memory' in row.keys():
                actual_peak = int(row['actual_peak_memory'])
                actual_memory[start] = actual_peak
            if start == -1:
                continue
            real_memory_changes[start] += size  # 在start时增加内存
            real_memory_changes[end - 1] += 0  # for plt
            real_memory_changes[end] -= size  # 在end时减少内存
        # theoretical memory
        for row in objects:
            producer = row['producer_task']
            if len(producer) == 0:
                continue
            producer_id = int(producer)
            input_str = row['user_tasks']
            type = row['type']
            # 使用正则表达式找到所有数字
            numbers = re.findall(r'\d+', input_str)
            # 将找到的数字字符串转换为整数列表
            user_tasks = [int(num) for num in numbers]
            if len(user_tasks) == 0:
                continue
            start = int(row['start_time_stamp'])
            end = int(row['end_time_stamp'])
            if type == 'SomasOutput' and start == -1:
                continue
            max_user_id = max(user_tasks)
            if end != -1 and end <= max_user_id:
                print(
                    f'Error: end time {end} is less than max user id {max_user_id}, producer_id {producer_id}, '
                    f'may be wrong ref count!!!!!')
            size = int(row['size'])
            theory_memory_changes[producer_id] += size  # 在start时增加内存
            theory_memory_changes[max_user_id] += 0  # for plt
            theory_memory_changes[max_user_id + 1] -= size  # 在end时减少内存

    # 画出内存随时间变化的折线图
    plt.figure(figsize=(40, 10))
    real_peak_time = plt_memory(real_memory_changes, 'memory_without_fragment')
    if boost:
        _ = plt_memory(theory_memory_changes, 'theory')
    _ = plt_memory(actual_memory, 'memory_with_fragment', False)
    plt.xlabel('Time')
    plt.ylabel('Memory Usage')
    plt.title('Memory Usage Over Time')
    plt.grid(True)
    plt.legend()
    plt.savefig('memory.png')

    # 把对应的峰值row存两个峰值csv
    reak_peak = []
    abnormal_mem = []
    for row in objects:
        start = int(row['start_time_stamp'])
        end = int(row['end_time_stamp'])
        if start == -1:
            continue
        if start <= real_peak_time <= end:
            reak_peak.append(row)
            input_str = row['user_tasks']
            type = row['type']
            # 使用正则表达式找到所有数字
            numbers = re.findall(r'\d+', input_str)
            # 将找到的数字字符串转换为整数列表
            user_tasks = [int(num) for num in numbers]
            if len(user_tasks) == 0:
                continue
            start = int(row['start_time_stamp'])
            end = int(row['end_time_stamp'])
            if type == 'SomasOutput' and start == -1:
                continue
            max_user_id = max(user_tasks)
            task_type = row['type']
            if max_user_id < real_peak_time and task_type.startswith("PyNative"):
                abnormal_mem.append(row)

    with open('peak.csv', mode='w', encoding='utf-8') as file:
        writer = csv.DictWriter(file, fieldnames=objects[0].keys())
        writer.writeheader()
        writer.writerows(reak_peak)

    if boost:
        with open('abnormal_mem.csv', mode='w', encoding='utf-8') as file:
            writer = csv.DictWriter(file, fieldnames=objects[0].keys())
            writer.writeheader()
            writer.writerows(abnormal_mem)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="ms runtime显存分析")
    parser.add_argument("block_file", type=str, help="block文件路径")
    parser.add_argument("--boost", action="store_true", help="是否使用增强分析功能")
    args = parser.parse_args()

    file_name = args.block_file
    print("file_name: ", file_name)
    tracker(file_name, args.boost)
